'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2024-09-23
Copyright © Department of Physics, Tsinghua University. All rights reserved

Add missing keys to old arXiv data
'''

import literature_format as lf
import pickle


def check_missing(item_dict: dict) -> bool:
    ''' Assuming the dict have the following keys:
            title, url
        If the dict lacks the following keys:
            author, date, abstract
        return False
    '''
    target_keys = ["author", "date", "abstract"]
    for curr_key in target_keys:
        if curr_key not in item_dict.keys():
            return False
    return True


def main_process_old_data(fname: str):
    ''' Process old data '''
    # 1. Get all data
    all_items = lf.markdown.parse_markdown_file(fname, is_old_format=True)
    print("Total:", len(all_items))

    # 2. Check missing
    new_items = []
    broken_items = []
    for curr_item in all_items:
        if 'url' not in curr_item.keys():
            print("Error: no url")
            print(curr_item)
            return
        if check_missing(curr_item):
            # Completed
            lf.arXiv.get_full_info(curr_item)
            new_items.append(curr_item)
        else:
            broken_items.append(curr_item)

    # 3. Get missing data
    repaired_items = []
    for start_ind in range(0, len(broken_items), 10):
        repaired_items += lf.arXiv.get_missing_info_from_api(
            broken_items[start_ind: start_ind + 10]
        )
    new_items += repaired_items

    # 4. Save
    print("Repaired:", len(repaired_items))
    print("New:", len(new_items))
    md_str = lf.markdown.save_to_markdown(new_items)
    with open("%s.processed.md" % (fname), "w", encoding='utf-8') as fp:
        print(md_str, file=fp)


def main_check_processed_data():
    with open("data/temp-files.pkl", "rb") as fp:
        all_processed = pickle.load(fp)

    for item in all_processed:
        all_keys = list(item.keys())
        all_keys.sort(key=lambda curr_key: lf.ALL_FIELDS.index(curr_key))
        print(all_keys)


if __name__ == '__main__':
    #### Tests ####
    def test_missing_check():
        all_items = lf.markdown.parse_markdown_file(r"data\arXiv\selected\D.Atomic_molecular_and_optical_physics.md", is_old_format=True)
        for item in all_items:
            print(check_missing(item))
            print(item.keys())
            input()

    def test_arXiv_query():
        with open("data/temp.xml", 'rb') as fp:
            response_content = fp.read()

        all_items = lf.arXiv.parse_arXiv_query(response_content)
        for item in all_items:
            print(item)

    #### Run ####
    main_process_old_data(
        r"data\arXiv\E.Optics.md"
    )
    main_process_old_data(
        r"data\arXiv\H.Surface_and_low-dimensional_physics.md"
    )
    main_process_old_data(
        r"data\arXiv\G.Nanoscopic_and_mesoscopic_physics.md"
    )
    main_process_old_data(
        r"data\arXiv\M.Quantum_information.md"
    )
    main_process_old_data(
        r"data\arXiv\N.Computational_physics.md"
    )
    main_process_old_data(
        r"data\arXiv\O.Statistical_physics_and_Complex_systems.md"
    )
